package com.it.controller;

import com.it.common.BaseContext;
import com.it.service.ChatMemoryAssistant;
import com.it.service.McpAssistant;

import dev.langchain4j.community.model.dashscope.WanxImageModel;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.mcp.client.DefaultMcpClient;
import dev.langchain4j.mcp.client.transport.McpTransport;
import dev.langchain4j.mcp.client.transport.stdio.StdioMcpTransport;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.ToolProvider;
import io.qdrant.client.QdrantClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Base64;
import java.util.List;
import java.util.Map;

@RestController
public class aiController {
    @Autowired
    private StreamingChatModel chatModelDeepSeek;
    @Autowired
    private ChatModel chatModelQWen;
    @Autowired
    private WanxImageModel wanxImageModel;
    @Autowired
    private StreamingChatModel chatModelQWenLong;
    @Autowired
    private ChatMemoryAssistant chatMemoryAssistant;
    @Autowired
    private ChatMemoryAssistant DeepSeekMemoryAssistant;
    @Autowired
    private QdrantClient qdrantClient;
    @PostMapping("/normal/chat")
    public Flux<String> normalChat(@RequestBody Map<String, Object> params) {
        String query = (String) params.get("query");
        return chatMemoryAssistant.chatWithMemory(BaseContext.getCurrentId(), query);
    }
    @PostMapping("/deepThinking/deepSeek")
    public Flux<String> deepSeek(@RequestBody Map<String, Object> params) {
        String query = (String) params.get("query");
        return DeepSeekMemoryAssistant.DeepSeekMemory(BaseContext.getCurrentId(), "AI回答："+query);
    }
    @PostMapping("/mcp/chat")
    public Flux<String> mcpChat(@RequestBody Map<String, Object> params)
    {
//        var vectorParams = Collections.VectorParams.newBuilder()
//                .setDistance(Collections.Distance.Cosine)
//                .setSize(1024)
//                .build();
//        qdrantClient.createCollectionAsync("test-qdrant",vectorParams);

        String query = (String) params.get("query");
        McpTransport builder = new StdioMcpTransport.Builder()
                .command(List.of("cmd", "/c","npx","-y","@baidumap/mcp-server-baidu-map"))
                .environment(Map.of("BAIDU_MAP_API_KEY", "9kDqxAG16mbEdgbja247Vyt1v9Q3DQiS"))
                .build();
        //构建McpClient客户端
        DefaultMcpClient build = new DefaultMcpClient.Builder()
                .transport(builder)
                .build();
        //创建工具集和原生的FunctionCalling模型
        ToolProvider toolProvider = McpToolProvider.builder().mcpClients(build).build();
        //通过AiService给我们自定义接口McpService构建实现类并将工具集和大模型赋值给AiService
        McpAssistant mcpService = AiServices.builder(McpAssistant.class)
                .streamingChatModel(chatModelDeepSeek)
                .toolProvider(toolProvider)
                .build();
        //调用大模型
        return mcpService.chat(query);
    }
    @PostMapping("/imageToText/chat")
    public String imageToText(@RequestParam("image") MultipartFile image,
                                    @RequestParam(value = "message", required = false) String message) throws IOException {
        byte[] byteArray = image.getBytes();
        String base64Image = Base64.getEncoder().encodeToString(byteArray);
        String mimeType = image.getContentType();
        String finalMessage = (message == null || message.isBlank()) ? "描述图片内容" : message;
        UserMessage userMessage = UserMessage.from(
                TextContent.from(finalMessage),
                ImageContent.from(base64Image, mimeType)
        );
        ChatResponse chat = chatModelQWen.chat(userMessage);
        String text = chat.aiMessage().text();
        return text;

    }
    @PostMapping("/textToImage/chat")
    public String imageToText(@RequestBody Map<String, Object> params) throws IOException {
        String query = (String) params.get("query");
        Response<Image> generate = wanxImageModel.generate(query);
        return generate.content().url().toString();
    }
}
